import os

import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import Dataset


class MedicalDataset(Dataset):
    def __init__(self, root_dir, is_train=True, image_size=256):
        """
        Args:
            root_dir (string): 根目录路径，包含train_images和test_images文件夹
            is_train (bool): 是否为训练模式
            image_size (int): 图像调整大小
        """
        self.root_dir = root_dir
        self.is_train = is_train
        self.image_size = image_size
        self.image_paths = []
        self.mask_paths = []

        # 图像转换
        self.image_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])

        # 根据是否训练模式选择相应文件夹
        if self.is_train:
            train_dir = os.path.join(root_dir, 'train_images')
            for folder in os.listdir(train_dir):
                folder_path = os.path.join(train_dir, folder)
                if os.path.isdir(folder_path):
                    img_folder = os.path.join(folder_path, 'images')
                    mask_folder = os.path.join(folder_path, 'masks')

                    # 检查文件夹是否存在
                    if not os.path.exists(img_folder) or not os.path.exists(mask_folder):
                        continue

                    img_files = os.listdir(img_folder)
                    mask_files = os.listdir(mask_folder)

                    if img_files and mask_files:  # 确保文件夹不为空
                        self.image_paths.append(os.path.join(img_folder, img_files[0]))
                        self.mask_paths.append(os.path.join(mask_folder, mask_files[0]))
        else:
            test_dir = os.path.join(root_dir, 'test_images')
            for folder in os.listdir(test_dir):
                folder_path = os.path.join(test_dir, folder)
                if os.path.isdir(folder_path):
                    test_files = os.listdir(folder_path)
                    if test_files:  # 确保文件夹不为空
                        self.image_paths.append(os.path.join(folder_path, test_files[0]))
                        self.mask_paths.append(None)

    def __len__(self):
        return len(self.image_paths)

    def preprocess_image(self, image):
        """预处理图像"""
        # 转换为PIL Image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        # 调整大小
        image = image.resize((self.image_size, self.image_size), Image.BILINEAR)
        # 应用转换
        image = self.image_transform(image)
        return image

    def preprocess_mask(self, mask):
        """预处理mask"""
        # 转换为PIL Image
        if isinstance(mask, np.ndarray):
            mask = Image.fromarray(mask)
        # 调整大小
        mask = mask.resize((self.image_size, self.image_size), Image.NEAREST)
        # 转换为tensor
        mask = torch.from_numpy(np.array(mask)).float()
        mask = (mask > 0).float()  # 二值化
        return mask

    def __getitem__(self, idx):
        # 读取图片
        img_path = self.image_paths[idx]
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = self.preprocess_image(image)

        if self.is_train:
            # 读取和处理mask
            mask_path = self.mask_paths[idx]
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            mask = self.preprocess_mask(mask)
            return image, mask
        else:
            return image


# 使用示例
if __name__ == '__main__':
    # 创建训练数据集
    train_dataset = MedicalDataset(
        root_dir='/Volumes/For_Mac/dateset/Pulmonary_X_ray_and_masks',
        is_train=True,
        image_size=256
    )

    # 创建测试数据集
    test_dataset = MedicalDataset(
        root_dir='/Volumes/For_Mac/dateset/Pulmonary_X_ray_and_masks',
        is_train=False,
        image_size=256
    )

    # 创建数据加载器
    from torch.utils.data import DataLoader

    train_loader = DataLoader(
        train_dataset,
        batch_size=4,
        shuffle=True,
        num_workers=4
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=1,
        shuffle=False,
        num_workers=4
    )

    # 检查数据
    for images, masks in train_loader:
        print("Training batch - Images:", images.shape, "Masks:", masks.shape)
        break

    for images in test_loader:
        print("Testing batch - Images:", images.shape)
        break
